#pragma once

#include <tuple>
#include <type_traits>

#include "result.h"

namespace limbo {

namespace details {

template <typename T, typename Context>
constexpr auto has_prefix = []() -> bool {
  if constexpr (std::is_same_v<
                    decltype(T::prefix(
                        std::declval<std::add_lvalue_reference_t<Context>>())),
                    uint32_t>) {
    return true;
  } else {
    return false;
  }
};

template <typename Layer, typename TargetState, size_t I>
constexpr size_t get_index_impl() {
  static_assert(I < std::tuple_size_v<Layer>, "no state on the layer");
  using State = std::tuple_element_t<I, Layer>;
  if constexpr (std::is_same_v<TargetState, State>) {
    return I;
  } else {
    return get_index_impl<Layer, TargetState, I + 1>();
  }
}

template <typename Layer, typename State> constexpr size_t get_index() {
  return get_index_impl<Layer, State, 0>();
}

template <typename Layer, typename TargetContext, size_t I>
constexpr size_t get_index_by_context_impl() {
  static_assert(I < std::tuple_size_v<Layer>, "no state on the layer");
  using State = std::tuple_element_t<I, Layer>;
  using Context = typename State::Context;
  if constexpr (std::is_same_v<TargetContext, Context>) {
    return I;
  } else {
    return get_index_by_context_impl<Layer, TargetContext, I + 1>();
  }
}

template <typename Layer, typename Context>
constexpr size_t get_index_by_context() {
  return get_index_by_context_impl<Layer, Context, 0>();
}

template <typename, typename = void> constexpr bool is_packetable{};

template <typename T>
constexpr bool is_packetable<T, std::void_t<typename T::Packet>> = true;

template <typename T> static constexpr auto get_prefixable_context(T &ctx) {
  using State = typename T::ContextState;
  auto lower_ctx = get_lower_context(ctx);
  if constexpr (!is_packetable<State>) {
    return get_prefixable_context(*lower_ctx);
  } else {
    return lower_ctx;
  }
}

template <typename, typename = void> constexpr bool is_frameable{true};

template <typename T>
constexpr bool is_frameable<T, std::void_t<typename T::non_frameable>> =
    T::non_frameable::value;

} // namespace details

template <typename... Ts> using Layer = std::tuple<Ts...>;

template <typename... Ls> struct Stack {
  using Layers = std::tuple<Ls...>;
  static constexpr auto N = std::tuple_size_v<Layers>;
  using LowestLayer = std::tuple_element_t<N - 1, Layers>;
  using LowestState = std::tuple_element_t<0, LowestLayer>;
  using EntryState = typename LowestState::LowerState;
  static_assert(std::tuple_size_v<LowestLayer> == 1,
                "lowest layer have to be only one");

  static_assert(std::is_default_constructible_v<Layers>,
                "every state should be default constructible");

  template <size_t I, size_t J> auto &get() {
    auto &layer = std::get<I>(layers);
    auto &state = std::get<J>(layer);
    return state;
  }

  template <typename Context>
  Result send(Chunk buff, Context &ctx, Chunk payload) {
    using Layer = std::tuple_element_t<0, Layers>;
    static constexpr auto J = details::get_index_by_context<Layer, Context>();
    auto prefixable_ctx = details::get_prefixable_context(ctx);
    auto prefix = prefix_impl(*prefixable_ctx);
    auto sub_buff = buff.skip(prefix);
    return send_impl<0, J, Context>(sub_buff, ctx, payload);
  }

  Result recv(Chunk buff, EntryState *entry_state) {
    auto &layer = std::get<N - 1>(layers);
    auto &state = std::get<0>(layer);
    auto result = state.recv(buff, entry_state);
    if (!result) {
      return result;
    }
    if (result.demand()) {
      return result;
    }
    if constexpr (N > 1) {
      auto inner_buff = state.get_parsed().payload;
      return recv_upper<N - 2, 0>(inner_buff, &state, result);
    }
    return result;
  }

private:
  template <size_t I, size_t J, typename Context>
  Result send_impl(Chunk buff, Context &ctx, Chunk payload) {
    using Layer = std::tuple_element_t<I, Layers>;
    using State = std::tuple_element_t<J, Layer>;
    if constexpr (details::is_packetable<State>) {
      using Packet = typename State::Packet;
      auto &layer = std::get<I>(layers);
      auto &state = std::get<J>(layer);
      auto result = Packet::send(state, ctx, buff, payload);
      if (!result) {
        return result;
      }
      if (result.demand()) {
        return result;
      }
      if constexpr (I + 1 < N) {
        using LowerLayer = std::tuple_element_t<I + 1, Layers>;
        using LowerState = typename State::LowerState;
        static constexpr auto K = details::get_index<LowerLayer, LowerState>();
        auto new_payload = result.consumed();
        auto lower_context = get_lower_context(ctx);
        auto prefix = get_prefix(*lower_context);
        auto sub_buff = buff.skip(-prefix);
        return send_impl<I + 1, K>(sub_buff, *lower_context, new_payload);
      } else {
        return result;
      }
    } else {
      using LowerLayer = std::tuple_element_t<I + 1, Layers>;
      using LowerState = typename State::LowerState;
      static constexpr auto K = details::get_index<LowerLayer, LowerState>();
      auto lower_ctx = get_lower_context(ctx);
      return send_impl<I + 1, K>(buff, *lower_ctx, payload);
    }
  }

  template <typename Context> uint32_t prefix_impl(const Context &ctx) {
    using State = typename Context::ContextState;
    auto value = get_prefix(ctx);
    if constexpr (details::has_lower_state_v<State>) {
      return value + prefix_impl(*get_lower_context(ctx));
    } else {
      return value;
    }
  }

  template <size_t I, size_t J, typename LowerState>
  Result recv_upper(Chunk buff, LowerState *lower_state, Result &lower_result) {
    using Layer = std::tuple_element_t<I, Layers>;
    using State = std::tuple_element_t<J, Layer>;
    if constexpr (!details::is_frameable<LowerState>) {
      if (!buff.size()) {
        return lower_result;
      }
    }
    if constexpr (std::is_same_v<typename State::LowerState, LowerState>) {
      auto &layer = std::get<I>(layers);
      auto &state = std::get<J>(layer);
      auto &parsed = lower_state->get_parsed();
      if (state.select(&parsed)) {
        auto result = state.recv(buff, lower_state);
        result.merge(lower_result.consumed());
        if constexpr (I) {
          if (!result || result.demand()) {
            return result;
          }
          auto inner_buff = state.get_parsed().payload;
          return recv_upper<I - 1, 0>(inner_buff, &state, result);
        }
        return result;
      }
    }
    if constexpr (J + 1 < std::tuple_size_v<Layer>) {
      return recv_upper<I, J + 1>(buff, lower_state, lower_result);
    }
    return lower_result;
  }

  Layers layers;
};

} // namespace limbo
